import time

import backtrader as bt
import pandas as pd
import datetime
import threading


class DynamicData(bt.feed.DataBase):
    params = (
        ('datetime', 0),
        ('open', 1),
        ('high', 2),
        ('low', 3),
        ('close', 4),
        ('volume', 5),
        ('openinterest', 6),
    )


class MyStrategy(bt.Strategy):
    params = (
        ("sma_period", 20),
    )

    def __init__(self):
        self.dataclose = self.data.close
        self.sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.sma_period)
        self.order = None

    def next(self):
        if self.dataclose > self.sma:
            # If the close is above the moving average, send a buy signal
            if self.order is None:
                self.order = self.buy()
        elif self.dataclose < self.sma:
            # If the close is below the moving average, send a sell signal
            if self.order is None:
                self.order = self.sell()

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"Buy: {order.executed.price}")
            elif order.issell():
                self.log(f"Sell: {order.executed.price}")

            self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f"Trade Profit: {trade.pnl}")


def data_gen():
    # Simulate real-time data generation
    for i in range(100):
        bar = {
            "datetime": datetime.datetime.now(),
            "open": 100 + i,
            "high": 105 + i,
            "low": 98 + i,
            "close": 102 + i,
            "volume": 1000
        }
        df = pd.DataFrame([bar], columns=["datetime", "open", "high", "low", "close", "volume"])

        # Create a custom data feed and add it to the cerebro
        data = DynamicData(dataname=df)
        cerebro.adddata(data)

        # Wait for a while to simulate real-time
        time.sleep(1)


if __name__ == "__main__":
    cerebro = bt.Cerebro()

    # Add the strategy
    cerebro.addstrategy(MyStrategy)

    # Set initial capital
    cerebro.broker.set_cash(100000)

    # Set commission
    cerebro.broker.setcommission(commission=0.001)

    print(f"Initial capital: {cerebro.broker.getvalue()}")

    # Start the data generation thread
    data_thread = threading.Thread(target=data_gen)
    data_thread.start()

    # Run the backtest
    cerebro.run()

    # Wait for the data generation thread to finish
    data_thread.join()

    print(f"Final capital: {cerebro.broker.getvalue()}")
